#include <iostream>
#include <time.h>
#include <thrust/device_vector.h>
#include <stdio.h>
#include <algorithm>
#include <vector>
#include "CudaRand.h"
using namespace std;

#define NUM_OBJECTS 200
#define LOOP_SIZE 2000
#define THREADS_PER_BLOCK 200
#define NUM_BLOCKS 300

#define INFECT_RATE 0.1

__host__ __device__
	void if_loop(RandGen& rand_obj){
		double temp = 0;
		double branch1 = rand_obj.cudaRandUni();
		if(branch1 < 0.3){
			temp *= 2.1;
		}else if(branch1 < 0.4 && branch1 >= 0.3){
			temp *= 3.1;
		}else{
			temp *= 4.1;
		}
		
		double branch2 = rand_obj.cudaRandUni();
		if(branch2 < 0.3){
			temp *= 2.1;
		}else if(branch2 < 0.4 && branch1 >= 0.3){
			temp *= 3.1;
		}else{
			temp *= 4.1;
		}
	}



__global__ void IfBranch(RandGen* rand_vec){

		RandGen& rand_obj = rand_vec[threadIdx.x+blockIdx.x*blockDim.x];
		int outer_size = rand_obj.cudaRandInt() % LOOP_SIZE;
		int inner_size = rand_obj.cudaRandInt() % LOOP_SIZE;
		
		if(rand_obj.cudaRandUni() < 0.1){
		
		for(int i = 0; i < outer_size; i++){
			if(rand_obj.cudaRandUni() < 0.1){
				for(int j = 0; j < inner_size; j++){
					if_loop(rand_obj);
				}
			}
			
			for(int j = 0; j < inner_size; j++){
				if(rand_obj.cudaRandUni() < 0.1){
					if_loop(rand_obj);
				}
			}

			
			for(int j = 0; j < inner_size; j++){
				if(rand_obj.cudaRandUni() < 0.1){
					if_loop(rand_obj);
				}
			}
			
		}
		
		}
	}

struct CPUBranch{
	void operator()(RandGen& rand_obj){
		int outer_size = rand_obj.cudaRandInt() % LOOP_SIZE;
		int inner_size = rand_obj.cudaRandInt() % LOOP_SIZE;
		
		if(rand_obj.cudaRandUni() < 0.1){
		
		for(int i = 0; i < outer_size; i++){
			if(rand_obj.cudaRandUni() < 0.1){
				for(int j = 0; j < inner_size; j++){
					if_loop(rand_obj);
				}
			}
			
			for(int j = 0; j < inner_size; j++){
				if(rand_obj.cudaRandUni() < 0.1){
					if_loop(rand_obj);
				}
			}

			
			for(int j = 0; j < inner_size; j++){
				if(rand_obj.cudaRandUni() < 0.1){
					if_loop(rand_obj);
				}
			}
			
		}
		
		}
	}
};	
	
	
__global__ void OptBranch(RandGen* rand_vec){
	__shared__ int remap_idx[THREADS_PER_BLOCK];
	__shared__ bool branch_idx[THREADS_PER_BLOCK];
	__shared__ bool temp_arr[THREADS_PER_BLOCK];
	
	int tid = threadIdx.x; 
	remap_idx[tid] = tid;
	RandGen& rand_obj = rand_vec[tid+blockIdx.x*blockDim.x];
	branch_idx[tid] = rand_obj.cudaRandUni() < 0.1;
	__syncthreads();
	

	int iter = 0;
	int reminder = THREADS_PER_BLOCK;
	while(reminder != 0){
		iter++;
		reminder = reminder >> 1;
	}
	int power = 1;
	
	for(int i = 0; i < iter; i++){
		// if(tid == 0 && blockIdx.x == 0){
			// for(int i = 0 ; i < THREADS_PER_BLOCK; i++){
				// printf("%d ",remap_idx[i]);
			// }
			// printf("\n");
			// for(int i = 0; i < THREADS_PER_BLOCK; i++){
				// printf("%d ",branch_idx[i]);
			// }
			// printf("\n\n");
		// }
		// __syncthreads();
	
		power *= 2;

		if(tid % power == 0){
			int p1 = tid;
			int p2 = tid + power / 2;
			int p = tid;

			while(true){
				if((p1 == (tid + power/2)) && ((p2 == tid + power) || p2 == THREADS_PER_BLOCK)) break;
				if(p1 >= (tid + power/2)){
					for(; ((p2 < (tid + power)) && p2 < THREADS_PER_BLOCK); p2++){
						temp_arr[p] = branch_idx[p2];
						p++;
					}
					break;
				}
				if(((p2 >= (tid + power)) || p2 >= THREADS_PER_BLOCK)){
					for(;(p1 < (tid + power/2)); p1++){
						temp_arr[p] = branch_idx[p1];
						p++;
					}
					break;
				}
				if(branch_idx[p1] > branch_idx[p2]){
					temp_arr[p] = branch_idx[p1];
					p1++;
				}
				else{
					temp_arr[p] = branch_idx[p2];
					p2++;
				}
				p++;
			}
		}
		__syncthreads();
		branch_idx[tid] = temp_arr[tid];
		__syncthreads();
	}

	int outer_size = rand_obj.cudaRandInt() % LOOP_SIZE;
	int inner_size = rand_obj.cudaRandInt() % LOOP_SIZE;
	if(branch_idx[tid]){
		
		for(int i = 0; i < outer_size; i++){
			if(rand_obj.cudaRandUni() < 0.1){
				for(int j = 0; j < inner_size; j++){
					if_loop(rand_obj);
				}
			}
			
			for(int j = 0; j < inner_size; j++){
				if(rand_obj.cudaRandUni() < 0.1){
					if_loop(rand_obj);
				}
			}

			
			for(int j = 0; j < inner_size; j++){
				if(rand_obj.cudaRandUni() < 0.1){
					if_loop(rand_obj);
				}
			}
			
		}
		
	}
	
}


/*
initial_status, end_status and rand_obj have the size of the entire population
start_idx, end_idx, thread_size has the size of number of blocks (or number of locations)
thread_size means how many people held by a thread
*/
__global__ void InfectionGPU(int* initial_status, int* end_status, int* start_idx, int* end_idx, RandGen* rand_vec,int* thread_size){
	int tid = threadIdx.x;
	int bid = blockIdx.x;

	int block_start_idx = start_idx[bid];
	int block_end_idx = end_idx[bid];

	int thread_len = thread_size[bid];
	int local_start_idx = block_start_idx  + tid * thread_len;
	int local_end_idx = local_start_idx + thread_len;
	local_end_idx = local_end_idx < block_end_idx ? local_end_idx : block_end_idx;
	
	__shared__ int shared_start_idx[THREADS_PER_BLOCK];
	__shared__ int shared_end_idx[THREADS_PER_BLOCK];

	shared_start_idx[tid] = local_start_idx;
	shared_end_idx[tid] = local_end_idx;
	
	__syncthreads();

	RandGen& rand_obj = rand_vec[bid];

	for(int i = local_start_idx; i < local_end_idx; i++){
		if(initial_status[i] == 1){
			for(int j = local_start_idx; j < local_end_idx; j++){
				if(initial_status[j] == 0 && rand_obj.cudaRandUni() < INFECT_RATE)
					end_status[j] = 1;
			}
		}
	}

	for(int i = 1; i < blockDim.x; i++){
		// forward infection
		int tmp_tid = tid + i;
		int tmp_start_idx = -1;
		int tmp_end_idx = -1;
		if(tmp_tid < blockDim.x){
			tmp_start_idx = shared_start_idx[tmp_tid];
			tmp_end_idx = shared_end_idx[tmp_tid];
			
			for(int i = local_start_idx; i < local_end_idx; i++){
				if(initial_status[i] == 1){
					for(int j = tmp_start_idx; j < tmp_end_idx; j++){
						if(initial_status[j] == 0 && rand_obj.cudaRandUni() < INFECT_RATE)
							end_status[j] = 1;
					}
				}
			}
		}

		tmp_tid = tid - i;
		if(tmp_tid >= 0){
			tmp_start_idx = shared_start_idx[tmp_tid];
			tmp_end_idx = shared_end_idx[tmp_tid];
			
			for(int i = local_start_idx; i < local_end_idx; i++){
				if(initial_status[i] == 1){
					for(int j = tmp_start_idx; j < tmp_end_idx; j++){
						if(initial_status[j] == 0 && rand_obj.cudaRandUni() < INFECT_RATE)
							end_status[j] = 1;
					}
				}
			}
		}
	}
}


__global__ void InfectionSimple(int* initial_status, int* end_status, int* start_idx, int* end_idx, RandGen* rand_vec){
	int tid = threadIdx.x;
	int bid = blockIdx.x;
	int idx = tid + bid * blockDim.x;
	
	int local_start_idx = start_idx[idx];
	int local_end_idx = end_idx[idx];
	
	RandGen& rand_obj = rand_vec[idx];
	
	for(int i = local_start_idx; i < local_end_idx; i++){
		if(initial_status[i] == 1){
			for(int j = local_start_idx; j < local_end_idx; j++){
				if(initial_status[j] == 0 && rand_obj.cudaRandUni() < INFECT_RATE)
					end_status[j] = 1;
			}
		}
	}
}

int main(){
	srand(time(NULL));
	cudaError err;
	
	int pp_size = 0; // size of population
	thrust::host_vector<int> h_start_idx(NUM_BLOCKS);
	thrust::host_vector<int> h_end_idx(NUM_BLOCKS);
	thrust::host_vector<int> h_thread_size(NUM_BLOCKS);
	for(int i = 0; i < NUM_BLOCKS; i++){
		h_start_idx[i] = pp_size;
		int comm_size = (10000 - rand()%5000);
		h_thread_size[i] = comm_size / THREADS_PER_BLOCK;
		pp_size += comm_size;
		h_end_idx[i] = pp_size;
	}
	thrust::host_vector<int> h_status(pp_size);
	// initial infection
	for(int i = 0; i < 30; i++) h_status[rand()%pp_size] = 1;
	thrust::host_vector<RandGen> h_rand_vec;
	for(int i = 0; i < pp_size; i++){
		h_rand_vec.push_back(RandGen(rand(),rand(),rand(),rand()));
	}
	
	thrust::device_vector<int> status = h_status;
	thrust::device_vector<int> start_idx = h_start_idx;
	thrust::device_vector<int> end_idx = h_end_idx;
	thrust::device_vector<int> thread_size = h_thread_size;
	thrust::device_vector<RandGen> rand_vec = h_rand_vec;
	thrust::device_vector<int> result(pp_size,0);
	
	int initial_infection = thrust::reduce(status.begin(),status.end());
	cout << "num of infection: " << initial_infection << endl;
	
	time_t start = clock();
	InfectionGPU<<<NUM_BLOCKS,THREADS_PER_BLOCK>>>(
		thrust::raw_pointer_cast(status.data()), 
		thrust::raw_pointer_cast(result.data()), 
		thrust::raw_pointer_cast(start_idx.data()),
		thrust::raw_pointer_cast(end_idx.data()),
		thrust::raw_pointer_cast(rand_vec.data()),
		thrust::raw_pointer_cast(thread_size.data()));
	
	err = cudaThreadSynchronize();
	if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	cout << "gpu1 which takes time: " << ((double)clock()-start)/CLOCKS_PER_SEC << endl;
	cout << "infected people: " << thrust::reduce(result.begin(),result.end()) << endl;
	
	
	thrust::device_vector<int> result2(pp_size,0);
	start = clock();
	InfectionSimple<<<1,NUM_BLOCKS>>>(
		thrust::raw_pointer_cast(status.data()), 
		thrust::raw_pointer_cast(result2.data()), 
		thrust::raw_pointer_cast(start_idx.data()),
		thrust::raw_pointer_cast(end_idx.data()),
		thrust::raw_pointer_cast(rand_vec.data()));
	
	err = cudaThreadSynchronize();
	if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	
	cout << "gpu2 which takes time: " << ((double)clock()-start)/CLOCKS_PER_SEC << endl;
	cout << "infected people (simple): " << thrust::reduce(result2.begin(),result2.end()) << endl;
	
	// thrust::device_vector<RandGen> rand_vec;
	// vector<RandGen> h_rand_vec;
	// for(int i = 0; i < NUM_OBJECTS; i++){
		// rand_vec.push_back(RandGen(rand(),rand(),rand(),rand()));
		// h_rand_vec.push_back(RandGen(rand(),rand(),rand(),rand()));
	// }
	
	// time_t start = clock();
	// IfBranch<<<NUM_BLOCKS,THREADS_PER_BLOCK>>>(thrust::raw_pointer_cast(rand_vec.data()));
	// err = cudaThreadSynchronize();
	// if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	// cout << "gpu1 which takes time: " << ((double)clock()-start)/CLOCKS_PER_SEC << endl;
	
	// start = clock();
	// for_each(h_rand_vec.begin(),h_rand_vec.end(),CPUBranch());
	// cout << "cpu which takes time: " << ((double)clock()-start)/CLOCKS_PER_SEC << endl;
	
	// start = clock();
	// OptBranch<<<NUM_BLOCKS,THREADS_PER_BLOCK>>>(thrust::raw_pointer_cast(rand_vec.data()));
	// err = cudaThreadSynchronize();
	// if(err != cudaSuccess) cout << "cudaThreadSynchronize error: " << cudaGetErrorString(err) << endl;
	// cout << "gpu2 which takes time: " << ((double)clock()-start)/CLOCKS_PER_SEC << endl;
	
	
	return 0;
}
